package raft

import (
	"context"
	"log"
	"math/rand"
	"sort"
	"sync"
	"sync/atomic"
	"time"
)

func (n *RaftNode) run() {
	ticker := time.NewTicker(10 * time.Millisecond)
	defer ticker.Stop()

	for {
		select {
		case <-n.stopChan:
			return
		case <-ticker.C:
			n.mu.Lock()
			switch n.state {
			case Follower:
				if time.Since(n.lastHeard) > n.electionTimeout {
					n.startElection()
				}
			case Candidate:
				if time.Since(n.lastHeard) > n.electionTimeout {
					n.startElection()
				}
			case Leader:
				n.broadcastAppendEntries()
				n.maybeTakeSnapshot()
			}
			n.mu.Unlock()
		}
	}
}

func (n *RaftNode) startElection() {
	n.state = Candidate
	n.currentTerm++
	n.votedFor = n.id
	n.persister.SaveState(n.currentTerm, n.votedFor, n.logs, n.lastSnapshotIndex, n.lastSnapshotTerm)

	votes := int32(1) // 投自己一票
	lastLogIndex := n.lastLogIndex()
	lastLogTerm := n.lastLogTerm()

	args := &RequestVoteReq{
		Term:         n.currentTerm,
		CandidateId:  n.id,
		LastLogIndex: lastLogIndex,
		LastLogTerm:  lastLogTerm,
	}

	var wg sync.WaitGroup

	for _, peer := range n.peers {
		wg.Add(1)
		go func(p *Peer) {
			defer wg.Done()
			ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
			defer cancel()

			resp, err := p.client.RequestVote(ctx, args)
			if err == nil && resp != nil && resp.VoteGranted {
				atomic.AddInt32(&votes, 1)
			}
		}(peer)
	}

	go func() {
		wg.Wait()
		n.mu.Lock()
		defer n.mu.Unlock()

		if n.state != Candidate {
			return
		}

		if votes > int32(len(n.peers)/2) {
			n.state = Leader
			n.initializeLeaderState()
			log.Printf("Node %d became leader for term %d", n.id, n.currentTerm)
		} else {
			n.state = Follower
		}
	}()

	// 重置选举超时
	n.lastHeard = time.Now()
	n.electionTimeout = time.Millisecond * time.Duration(150+rand.Intn(150))
}

func (n *RaftNode) broadcastAppendEntries() {
	for id, peer := range n.peers {
		if n.nextIndex[id] <= n.lastSnapshotIndex {
			n.sendSnapshot(id)
		} else {
			n.sendAppendEntries(id)
		}
	}
}

func (n *RaftNode) sendAppendEntries(peerID int32) {
	nextIdx := n.nextIndex[peerID]
	prevLogIndex := nextIdx - 1
	prevLogTerm := n.logTerm(prevLogIndex)

	var entries []*LogEntry
	if nextIdx <= n.lastLogIndex() {
		end := min(nextIdx+50, n.lastLogIndex())
		for i := nextIdx; i <= end; i++ {
			if entry, err := n.logAt(i); err == nil {
				entries = append(entries, entry)
			}
		}
	}

	args := &AppendEntriesReq{
		Term:         n.currentTerm,
		LeaderId:     n.id,
		PrevLogIndex: prevLogIndex,
		PrevLogTerm:  prevLogTerm,
		Entries:      entries,
		LeaderCommit: n.commitIndex,
	}

	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()

	resp, err := n.peers[peerID].client.AppendEntries(ctx, args)
	if err != nil {
		return
	}

	n.mu.Lock()
	defer n.mu.Unlock()

	if resp.Term > n.currentTerm {
		n.currentTerm = resp.Term
		n.state = Follower
		n.votedFor = -1
		n.persister.SaveState(n.currentTerm, n.votedFor, n.logs, n.lastSnapshotIndex, n.lastSnapshotTerm)
		return
	}

	if resp.Success {
		n.nextIndex[peerID] = nextIdx + int64(len(entries))
		n.matchIndex[peerID] = n.nextIndex[peerID] - 1
		n.advanceCommitIndex()
	} else {
		if resp.ConflictIndex > 0 {
			n.nextIndex[peerID] = resp.ConflictIndex
		} else {
			n.nextIndex[peerID] = max(1, nextIdx-1)
		}
	}
}

func (n *RaftNode) advanceCommitIndex() {
	// 找到多数节点已复制的最大索引
	indexes := []int64{}
	for _, idx := range n.matchIndex {
		indexes = append(indexes, idx)
	}
	sort.Slice(indexes, func(i, j int) bool { return indexes[i] < indexes[j] })
	majorityIndex := indexes[len(indexes)/2]

	if majorityIndex > n.commitIndex && n.logTerm(majorityIndex) == n.currentTerm {
		n.commitIndex = majorityIndex
	}
}
